import yaml


def parse_yaml_modules(file_path):
    """
    解析 YAML 配置文件，提取所有 module_name。

    :param file_path: YAML 配置文件路径
    :return: 列表，包含所有 module_name
    """
    with open(file_path, 'r') as stream:
        try:
            config = yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            print(f"Error parsing YAML: {exc}")
            return None

    modules = config.get('modules', {})
    module_list = []

    for module_category in modules.values():
        for module in module_category:
            if 'module_name' in module:
                module_list.append(module['module_name'])
    module_length = len(module_list)

    training_param = config.get('training_param', {})
    return training_param, module_list, module_length


# if __name__ == '__main__':
#     file_path = "/home/lionrock-g4/wangjiaqian/codes/code_test/sample_test/llm_simulator_by_alpa_layer/config/model_config_qwen3_32b.yaml"  # 修改为你的 YAML 文件路径
#     training_param, module_list, module_length = parse_yaml_modules(file_path)

#     print("Extracted module names:")
#     for name in module_list:
#         print(name)
#     print(module_list)
#     print(module_length)
